{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MCP Proxy\n",
    "\n",
    "The [Model Context Protocol (MCP)](https://modelcontextprotocol.io/introduction) is an open protocol that defines a standardized way to provide context to large language models (LLMs). By using MCP, applications can seamlessly interact with LLMs in a structured and efficient manner, making it easier to exchange information, use external tools, and enhance the capabilities of language models in a consistent and reusable way.\n",
    "\n",
    "This guide will walk you through how to create **MCP clients** that connect to MCP servers, enabling you to utilize external tools and processes efficiently.\n",
    "\n",
    "## Introduction to MCP\n",
    "\n",
    "MCP aims to standardize how applications provide contextual information to LLMs and how LLMs can access external tools. MCP servers expose specific tools that are accessible via different communication protocols, such as [`stdio` (Standard Input/Output)](https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#stdio), [`SSE` (Server-Sent Events)](https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#http-with-sse), and [`streamable-http`(Streamable HTTP)](https://modelcontextprotocol.io/specification/2025-06-18/basic/transports#streamable-http).\n",
    "\n",
    "MCP clients are responsible for interacting with these servers, sending requests, and processing responses. By integrating MCP into your application, you can build systems where LLMs can utilize external computational tools in real-time.\n",
    "\n",
    "In this guide, we will:\n",
    "\n",
    "- Set up an MCP server that exposes basic mathematical tools (`add` and `multiply`).\n",
    "- Connect the AG2 framework to this MCP server.\n",
    "- Communicate using two different transport protocols: [stdio](https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#stdio) and [SSE](https://spec.modelcontextprotocol.io/specification/2024-11-05/basic/transports/#http-with-sse).\n",
    "\n",
    "## Installation\n",
    "\n",
    "To integrate MCP tools into the AG2 framework, install the required dependencies:\n",
    "\n",
    "```bash\n",
    "pip install -U \"ag2[openai,mcp,mcp-proxy-gen]>=0.9.4\"\n",
    "```\n",
    "\n",
    "> **Note:** If you are using `pyautogen`, uninstall it before installing `ag2` to ensure compatibility:\n",
    "> ```bash\n",
    "> pip uninstall -y pyautogen\n",
    "> ```\n",
    "> ```bash\n",
    "> pip install -U \"ag2[openai,mcp,mcp-proxy-gen]>=0.9.4\"\n",
    "> ``` \n",
    "\n",
    "## Imports\n",
    "\n",
    "Before diving into the code, let’s go over the main imports used in this guide:\n",
    "\n",
    "- `ClientSession`: Manages the client session for connecting to the MCP server, allowing you to send requests and handle responses.\n",
    "- `StdioServerParameters`: Provides parameters to set up communication with an MCP server over `stdio`.\n",
    "- `stdio_client`: A utility that helps you connect to an MCP server using the standard input/output protocol.\n",
    "- `sse_client`: A utility to connect to an MCP server using the `SSE` (Server-Sent Events) protocol.\n",
    "- `create_toolkit`: A helper function that creates a toolkit, wrapping the tools exposed by the MCP server for easier access.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import timedelta\n",
    "from pathlib import Path\n",
    "\n",
    "from mcp import ClientSession, StdioServerParameters\n",
    "from mcp.client.sse import sse_client\n",
    "from mcp.client.stdio import stdio_client\n",
    "\n",
    "from autogen import LLMConfig\n",
    "from autogen.agentchat import AssistantAgent\n",
    "from autogen.mcp import create_toolkit\n",
    "from autogen.mcp.mcp_proxy import MCPProxy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Only needed for Jupyter notebooks\n",
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate Proxy MCP Server Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "\n",
    "mcp_whatsapp_path = Path() / \"mcp_whatsapp\"\n",
    "shutil.rmtree(mcp_whatsapp_path, ignore_errors=True)\n",
    "mcp_whatsapp_path.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# Create an MCP server code via the MCPProxy (requires dependencies from mcp-proxy-gen option)\n",
    "MCPProxy.create(\n",
    "    openapi_url=\"https://dev.infobip.com/openapi/products/whatsapp.json\",\n",
    "    client_source_path=mcp_whatsapp_path,\n",
    "    servers=[{\"url\": \"https://api.infobip.com\"}],\n",
    ")\n",
    "\n",
    "assert mcp_whatsapp_path.exists(), \"Failed to create tmp directory\"\n",
    "# assert list(tmp_path.glob(\"*.py\")) >= 2, \"Failed to generate OpenAPI client files\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start Proxy server"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting Up an MCP Server\n",
    "\n",
    "We will create a simple MCP server that exposes two tools: `add` and `multiply`.\n",
    "The server listens for requests using either the **stdio** or **SSE** transport protocols, depending on the argument passed when starting the server.\n",
    "\n",
    "For more details on creating MCP servers, visit the [MCP Python SDK documentation](https://github.com/modelcontextprotocol/python-sdk).\n",
    "\n",
    "Will be generated (`math_server.py`) with the following content:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def print_available_tools(session: ClientSession) -> None:\n",
    "    list_tools_result = await session.list_tools()\n",
    "    tools = list_tools_result.tools\n",
    "    print(\"Available tools:\")\n",
    "    for tool in tools:\n",
    "        print(f\"- {tool.name}: {tool.description} ({tool.inputSchema})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create server parameters for stdio connection\n",
    "\n",
    "\n",
    "server_params = StdioServerParameters(\n",
    "    command=\"python\",  # The command to run the server\n",
    "    args=[\n",
    "        str(mcp_whatsapp_path),\n",
    "        \"stdio\",\n",
    "    ],  # Path to server script and transport mode\n",
    ")\n",
    "\n",
    "\"\"\"Context manager for stdio client.\"\"\"\n",
    "async with (\n",
    "    stdio_client(server_params) as (read, write),\n",
    "    ClientSession(read, write, read_timeout_seconds=timedelta(seconds=30)) as session,\n",
    "):\n",
    "    # Initialize the connection\n",
    "    await session.initialize()\n",
    "    await print_available_tools(session)\n",
    "\n",
    "    # make a simple call to add function\n",
    "    params = {\"a\": 2, \"b\": 3}\n",
    "    add_result = await session.call_tool(\"add\", params)\n",
    "    print()\n",
    "    print(f\"Result of calling tool 'add' with params '{params}' is '{add_result.content[0].text}'\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating a Toolkit from MCP Tools\n",
    "\n",
    "To allow AG2 to interact with the MCP server, we need to establish a connection, create a toolkit from the tools exposed by the server, and then register this toolkit with an AG2 agent. The toolkit will provide the necessary functionality for AG2 to call the MCP tools, such as performing mathematical operations.\n",
    "\n",
    "**Steps to Create a Toolkit from MCP Tools**\n",
    "\n",
    "- **Wrap MCP tools into a toolkit**: The tools exposed by the MCP server (like `add` and `multiply`) are wrapped into a toolkit for easy use by AG2.\n",
    "- **Register the toolkit with an AG2 agent**: This step allows AG2 to use the toolkit’s tools when processing requests.\n",
    "- **Start a conversation**: AG2 can now send a message to the MCP server requesting a task to be performed (e.g., adding two numbers) and receive the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def create_toolkit_and_run(session: ClientSession) -> None:\n",
    "    # Create a toolkit with available MCP tools\n",
    "    toolkit = await create_toolkit(session=session)\n",
    "    agent = AssistantAgent(\n",
    "        name=\"assistant\", llm_config=LLMConfig(config_list={\"model\": \"gpt-5-nano\", \"api_type\": \"openai\"})\n",
    "    )\n",
    "    # Register MCP tools with the agent\n",
    "    toolkit.register_for_llm(agent)\n",
    "\n",
    "    # Make a request using the MCP tool\n",
    "    result = await agent.a_run(\n",
    "        message=\"Add 123223 and 456789\",\n",
    "        tools=toolkit.tools,\n",
    "        max_turns=2,\n",
    "        user_input=False,\n",
    "    )\n",
    "    await result.process()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting Up the Client Session and Communicating with the MCP Server\n",
    "\n",
    "Once the toolkit is created, the next step is to connect AG2 to the MCP server and manage the communication between them. This involves setting up connection parameters and starting a client session, which will facilitate the entire interaction. After establishing the connection, AG2 can send requests and receive responses from the server.\n",
    "\n",
    "### Steps to Set Up the Client Session and Communicate with the MCP Server\n",
    "\n",
    "- **Set up connection parameters**: Define how AG2 will communicate with the MCP server. This includes specifying the transport protocol (`stdio` or `SSE` or `streamable-http`) and other connection details.\n",
    "- **Start a client session**: This session facilitates communication between AG2 and the MCP server. The session is responsible for sending requests, receiving responses, and maintaining the connection state.\n",
    "\n",
    "#### Using `stdio_client` for Communication\n",
    "\n",
    "Establish the connection in the client code by using the `stdio_client` requires:\n",
    "\n",
    "- **Server Parameters**: The connection is defined using the `StdioServerParameters`, where the command (`python`) runs the server (`math_server.py`) with `stdio` mode.\n",
    "- **Connecting to the Server**: The `stdio_client` establishes the connection to the server, and the `ClientSession` is used to facilitate communication.\n",
    "- **Running the Task**: Once the connection is established, `create_toolkit_and_run(session)` is called to wrap the MCP tools and perform the task asynchronously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create server parameters for stdio connection\n",
    "server_params = StdioServerParameters(\n",
    "    # add security parameters here as env var\n",
    "    command=\"python\",  # The command to run the server\n",
    "    args=[\n",
    "        str(mcp_whatsapp_path / \"main_tmp.py\"),  # Path to the server script\n",
    "        \"stdio\",\n",
    "    ],  # Path to server script and transport mode\n",
    "    env={\n",
    "        # Add environment variables needed for security here\n",
    "    },\n",
    ")\n",
    "\n",
    "async with (\n",
    "    stdio_client(server_params) as (read, write),\n",
    "    ClientSession(read, write, read_timeout_seconds=timedelta(seconds=30)) as session,\n",
    "):\n",
    "    # Initialize the connection\n",
    "    await session.initialize()\n",
    "    await create_toolkit_and_run(session)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Using `sse_client` for Communication\n",
    "\n",
    "To interact with the MCP server via `SSE`, follow these steps:\n",
    "\n",
    "- Open new terminal and start the server using the sse transport mode:\n",
    "\n",
    "    > ```bash\n",
    "    > python math_server.py sse\n",
    "    > ```\n",
    "- Once the server is running, rstablish the connection in the client code by using the `sse_client`:\n",
    "  \n",
    "  - **Server URL**: The `sse_client` connects to the `SSE` server running locally at `http://127.0.0.1:8000/sse`.\n",
    "  - **Connecting to the Server**: The `sse_client` establishes an `SSE` connection, and the `ClientSession` is used to manage the communication with the server.\n",
    "  - **Running the Task**: Once the session is initialized, `create_toolkit_and_run(session)` is called to create the toolkit from MCP tools and perform the task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async with sse_client(url=\"http://127.0.0.1:8000/sse\") as streams, ClientSession(*streams) as session:\n",
    "    # Initialize the connection\n",
    "    await session.initialize()\n",
    "    await create_toolkit_and_run(session)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "from mcp import ClientSession, StdioServerParameters\n",
    "from mcp.client.sse import sse_client\n",
    "from mcp.client.stdio import stdio_client\n",
    "\n",
    "from autogen.mcp.mcp_proxy import MCPProxy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MCPProxy:\n",
    "    def __init__(self, openapi_url: str, client_source_path: Path, servers: list) -> None:\n",
    "        \"\"\"\n",
    "        Initialize the MCP Proxy server.\n",
    "\n",
    "        Args:\n",
    "            openapi_url (str): URL to the OpenAPI specification.\n",
    "            client_source_path (Path): Path to the generated client source.\n",
    "            servers (list): List of servers to connect to.\n",
    "        \"\"\"\n",
    "        self.openapi_url = openapi_url\n",
    "        self.client_source_path = client_source_path\n",
    "        self.servers = servers\n",
    "\n",
    "    @classmethod\n",
    "    def create(cls, openapi_url: str, client_source_path: Path, servers: list) -> None:\n",
    "        \"\"\"\n",
    "        Create an MCP Proxy server based on the OpenAPI specification.\n",
    "\n",
    "        Args:\n",
    "            openapi_url (str): URL to the OpenAPI specification.\n",
    "            client_source_path (Path): Path to the generated client source.\n",
    "            servers (list): List of servers to connect to.\n",
    "        \"\"\"\n",
    "        # Implementation of MCPProxy creation\n",
    "        return MCPProxy(\n",
    "            openapi_url=openapi_url,\n",
    "            client_source_path=client_source_path,\n",
    "            servers=servers,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client_source_path = Path(...)\n",
    "\n",
    "# Create an MCP Proxy server based on the OpenAPI specification\n",
    "MCPProxy.create(\n",
    "    openapi_url=\"https://dev.infobip.com/openapi/products/whatsapp.json\",  # OpenAPI spec URL\n",
    "    client_source_path=client_source_path,  # Path to the generated client source\n",
    "    servers=[{\"url\": \"https://api.infobip.com\"}],  # List of servers to connect to\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def create_toolkit_and_run(session: ClientSession) -> None:\n",
    "    # Create a toolkit with available MCP tools\n",
    "    toolkit = await create_toolkit(session=session)\n",
    "\n",
    "    # Initialize an AssistantAgent with GPT-4o-mini model\n",
    "    agent = AssistantAgent(\n",
    "        name=\"assistant\", llm_config=LLMConfig(config_list={\"model\": \"gpt-5-nano\", \"api_type\": \"openai\"})\n",
    "    )\n",
    "\n",
    "    # Register the toolkit with the agent\n",
    "    toolkit.register_for_llm(agent)\n",
    "\n",
    "    # Make a request using the MCP tool\n",
    "    result = await agent.a_run(\n",
    "        message=\"Add 123223 and 456789\",  # Example message for the tool\n",
    "        tools=toolkit.tools,\n",
    "        max_turns=2,\n",
    "        user_input=False,\n",
    "    )\n",
    "    await result.process()\n",
    "\n",
    "\n",
    "# Create server parameters for stdio connection, including environment variables for secrets\n",
    "server_params = StdioServerParameters(\n",
    "    command=\"python\",  # The command to run the server\n",
    "    args=[\n",
    "        str(mcp_whatsapp_path / \"main_tmp.py\"),  # Path to the server script\n",
    "        \"stdio\",  # Stdio transport mode\n",
    "    ],\n",
    "    env={\n",
    "        \"API_KEY\": \"your-api-key\",  # Example API Key passed as an environment variable\n",
    "        \"OAUTH2_TOKEN\": \"your-oauth2-token\",  # Example OAuth2 token passed as an environment variable\n",
    "    },\n",
    ")\n",
    "\n",
    "# Initialize the stdio connection and create the MCP server session\n",
    "async with (\n",
    "    stdio_client(server_params) as (read, write),\n",
    "    ClientSession(read, write, read_timeout_seconds=timedelta(seconds=30)) as session,\n",
    "):\n",
    "    await session.initialize()  # Initialize the session\n",
    "    await create_toolkit_and_run(session)  # Run the toolkit and agent"
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "MCP Clients",
   "tags": [
    "MCP",
    "Model Context Protocol",
    "tools"
   ]
  },
  "kernelspec": {
   "display_name": ".venv-3.11-openapi",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
